package look.word.datasource.config;

import org.springframework.context.annotation.Primary;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.sql.DataSource;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author : look-word
 * 2022-10-11 22:19
 **/
@Component
@Primary
public class DataSourceManagement extends AbstractRoutingDataSource {

    public static ThreadLocal<String> flag = new ThreadLocal<>();

    /**
     * 注入数据源
     */
    @Resource
    private DataSource mysqlDataSource1;
    /**
     * 注入数据源
     */
    @Resource
    private DataSource mysqlDataSource2;

    public DataSourceManagement() {
        flag.set(DataSourceType.MYSQL_DATASOURCE1.name());
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return flag.get();
    }

    @Override
    public void afterPropertiesSet() {
        Map<Object, Object> targetDataSource = new ConcurrentHashMap<>();
        targetDataSource.put(DataSourceType.MYSQL_DATASOURCE1.name(), mysqlDataSource1);
        targetDataSource.put(DataSourceType.MYSQL_DATASOURCE2.name(), mysqlDataSource2);
        // 设置数据源来源
        super.setTargetDataSources(targetDataSource);
        // 设置默认数据源
        super.setDefaultTargetDataSource(mysqlDataSource1);
        super.afterPropertiesSet();
    }
}
